import cv2
import os
import numpy as np
from tqdm import tqdm
import multiprocessing
from multiprocessing import Pool
from functools import partial

class ImageCropper:
    def __init__(self):
        self.crop_width = None  # 裁剪宽度
        self.reference_right_edge = None  # 参考右边缘位置
        
    def find_right_edge(self, binary_img):
        # 计算垂直投影
        projection = np.sum(binary_img, axis=0)

        # 使用一维高斯卷积减少噪声影响
        window = 15 # 窗口大小
        sigma = window/3.0
        x = np.linspace(-window//2, window//2, window)
        gaussian = np.exp(-(x**2)/(2*sigma**2))
        gaussian = gaussian/np.sum(gaussian)
        smoothed = np.convolve(projection, gaussian, mode='valid')

        # 计算梯度
        gradient = np.gradient(smoothed)

        # 从右向左扫描找到第一个显著边缘
        threshold = np.std(gradient) * 2
        for i in range(len(gradient)-1, window, -1):
            if abs(gradient[i]) > threshold:
                # 使用局部最大值优化边缘位置
                local_region = gradient[max(0, i-5):min(len(gradient), i+6)]
                local_offset = np.argmax(np.abs(local_region))
                return i - 5 + local_offset
        return None
    
    def initialize_with_reference(self, reference_img):
        # 转换为灰度图
        gray = cv2.cvtColor(reference_img, cv2.COLOR_BGR2GRAY)
        
        # 增强对比度
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到右边缘
        right_edge = self.find_right_edge(binary)
        if right_edge is not None:
            self.reference_right_edge = right_edge
            # 设置裁剪宽度（从右边缘向左的距离）
            self.crop_width = right_edge - 100  # 可以根据需要调整这个值
    
    def crop_image(self, img):
        if self.crop_width is None:
            return img
            
        # 转换为灰度图并增强对比度
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
        gray = clahe.apply(gray)
        
        # 检测垂直线
        binary = detect_vertical_lines(gray)
        
        # 找到当前帧的右边缘
        current_right_edge = self.find_right_edge(binary)
        
        if current_right_edge is None:
            return img
            
        # 计算裁剪区域，使用动态边缘位置
        start_x = max(0, current_right_edge - self.crop_width)
        end_x = current_right_edge
        
        # 裁剪图像
        cropped = img[:, start_x:end_x]
        
        return cropped

def detect_vertical_lines(gray_img):
    # 使用Sobel算子增强垂直边缘检测
    sobelx = cv2.Sobel(gray_img, cv2.CV_64F, 1, 0, ksize=5)
    sobelx = cv2.convertScaleAbs(sobelx)

    # 使用全局阈值
    _, binary = cv2.threshold(sobelx, 100, 255, cv2.THRESH_BINARY)

    # 形态学操作增强垂直线条
    kernel_vertical = np.ones((15,1), np.uint8)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel_vertical)
    binary = cv2.morphologyEx(binary, cv2.MORPH_OPEN, kernel_vertical)

    return binary

def process_single_image(args):
    img_path, cropper = args
    try:
        # 读取当前帧
        current_frame = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if current_frame is None:
            return None
        
        # 裁剪图像
        cropped_frame = cropper.crop_image(current_frame)
        
        return cropped_frame
    except Exception as e:
        print(f"处理图片 {img_path} 时发生错误: {str(e)}")
        return None

def process_images_to_video(input_path, output_path, sample_count=4000):
    try:
        # 获取所有bmp文件
        image_files = [f for f in os.listdir(input_path) if f.lower().endswith('.bmp')]
        if not image_files:
            print(f"错误：在 {input_path} 中没有找到BMP文件")
            return
        image_files.sort()
        
        total_images = len(image_files)
        if total_images < sample_count:
            sample_count = total_images
        step = total_images // sample_count
        
        # 读取第一张图片作为参考帧
        first_image_path = os.path.join(input_path, image_files[0])
        print(f"正在读取第一张图片: {first_image_path}")
        reference_frame = cv2.imdecode(np.fromfile(first_image_path, dtype=np.uint8), cv2.IMREAD_COLOR)
        
        if reference_frame is None:
            print(f"错误：无法读取图片 {first_image_path}")
            return
        
        # 创建裁剪器并初始化
        cropper = ImageCropper()
        cropper.initialize_with_reference(reference_frame)
        
        # 获取裁剪后的尺寸
        first_processed = cropper.crop_image(reference_frame)
        height, width = first_processed.shape[:2]
        
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        # 准备并行处理的参数
        selected_files = [os.path.join(input_path, image_files[i]) for i in range(0, total_images, step)]
        process_args = [(f, cropper) for f in selected_files]
        
        # 获取CPU核心数
        num_cores = multiprocessing.cpu_count()
        print(f"使用 {num_cores} 个CPU核心进行并行处理")
        
        # 创建进程池
        with Pool(num_cores) as pool:
            # 使用tqdm显示进度
            processed_frames = list(tqdm(
                pool.imap(process_single_image, process_args),
                total=len(process_args),
                desc="处理图片"
            ))
        
        # 过滤掉处理失败的帧
        processed_frames = [frame for frame in processed_frames if frame is not None]
        
        if not processed_frames:
            print("错误：没有成功处理任何图片")
            return
        
        # 创建视频写入器
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))
        
        # 写入处理后的帧
        for frame in processed_frames:
            out.write(frame)
        
        out.release()
        print("视频生成完成！")
        
    except Exception as e:
        print(f"发生错误: {str(e)}")

if __name__ == '__main__':
    # 使用原始字符串表示路径
    input_path = r"F:\李双成\B-1\200-4"  # 使用原始字符串
    output_path = r"F:\test\output_video_stabilized.mp4"  # 使用原始字符串
    process_images_to_video(input_path, output_path)
